from neuracle_lib.readbdfdata import readbdfdata
from tkinter import filedialog
from tkinter import *
import numpy as np
import os


def check_files_format(path):
    """
    检查文件格式并返回文件名和路径。

    参数：
    path (list of str): 选择的文件路径列表。

    返回：
    tuple: 包含文件名列表和路径列表的元组。
    """
    filename = []
    pathname = []
    if len(path) == 0:
        raise TypeError('please select valid file')

    elif len(path) == 1:
        (temppathname, tempfilename) = os.path.split(path[0])
        if 'edf' in tempfilename:
            filename.append(tempfilename)
            pathname.append(temppathname)
            return filename, pathname
        elif 'bdf' in tempfilename:
            raise TypeError('unsupport only one neuracle-bdf file')
        else:
            raise TypeError('not support such file format')

    else:
        temp = []
        temppathname = r''
        evtfile = []
        idx = np.zeros((len(path) - 1,))
        for i, ele in enumerate(path):
            (temppathname, tempfilename) = os.path.split(ele)
            if 'data' in tempfilename:
                temp.append(tempfilename)
                if len(tempfilename.split('.')) > 2:
                    try:
                        idx[i] = (int(tempfilename.split('.')[1]))
                    except:
                        raise TypeError('no such kind file')
                else:
                    idx[i] = 0
            elif 'evt' in tempfilename:
                evtfile.append(tempfilename)

        pathname.append(temppathname)
        datafile = [temp[i] for i in np.argsort(idx)]

        if len(evtfile) == 0:
            raise TypeError('not found evt.bdf file')

        if len(datafile) == 0:
            raise TypeError('not found data.bdf file')
        elif len(datafile) > 1:
            print('current readbdfdata() only support continue one data.bdf ')
            return filename, pathname
        else:
            filename.append(datafile[0])
            filename.append(evtfile[0])
            return filename, pathname


def extract_data():
    root = Tk()
    root.withdraw()
    # 弹窗选择脑电文件
    path = filedialog.askopenfilenames(title='Select two bdf files',
                                       filetypes=(("two bdf files", "*.bdf"), ("one edf files", "*.edf")))
    # 检查文件类型
    filename, pathname = check_files_format(path)
    ## 读取数据
    eeg = readbdfdata(filename, pathname)

    # 获取目标通道
    channels = ['C3','C4']
    channel_select = [eeg['ch_names'].index(i) for i in channels]
    # 将原始数据转换为矩阵并选取目标通道数据
    data_matrix = np.asarray(eeg['data'])[channel_select, :]

    # 保存数据和标签
    dataset = []
    labels = []
    duration = 1000
    for event in eeg['events']:
        # event存储的格式为[beginSamplePoint,Duration,Event]

        dataset.append(data_matrix[:, event[0]:event[0] + duration])
        labels.append(event[2])

    dataset = np.asarray(dataset)
    labels = np.asarray(labels)
    np.save('dataset0.npy', {'dataset': dataset, 'labels': labels})



def save_data():
    root = Tk()
    root.withdraw()
    # 弹窗选择脑电文件
    path = filedialog.askopenfilenames(title='Select two bdf files',
                                       filetypes=(("two bdf files", "*.bdf"), ("one edf files", "*.edf")))
    # 检查文件类型
    filename, pathname = check_files_format(path)
    ## 读取数据
    eeg = readbdfdata(filename, pathname)

    # 获取目标通道
    channels = eeg['ch_names']
    channel_select = [eeg['ch_names'].index(i) for i in channels]
    # 将原始数据转换为矩阵并选取目标通道数据
    data_matrix = np.asarray(eeg['data'])[channel_select, :]

    # 保存数据和标签
    np.savez('dataset0.npz', data = data_matrix,channels = eeg['ch_names'])


if __name__ == '__main__':
    save_data()
